# =========================================================================
# ||  This module includes some trick function. Such as match catalog.
# =========================================================================
from astropy.coordinates import SkyCoord
import astropy.units as u
import numpy as np

from astropy.coordinates import SkyCoord
import astropy.units as u
def match_ctg(target_coords,candidates_inds,candidates_coords, target_inds=None, 
    matchtype='best'):
	# INPUT PARAMETERS: 
	#       target_coords: Two elements array or array with shape (n,2)
    #       candidates_inds: Array, containing indices of m points.
	#       candidates_coords: Array, containing m points with 
	# 			coordinates of shape (m,2).
    #       matchtype: string, same as topcat
	# OUTPUT PARAMETERS:
	# 		output = {'target indices':output_inds_target_list,
    #       'matched indices':output_inds_cands_list,
    #       'ra_cands':output_ra_cands_list,
    #       'dec_cands':output_dec_cands_list,
    #       'ra_targs':output_ra_target_list,
    #       'dec_targs':output_dec_target_list,
    #       'separation':output_sep_list
    #       }
	# EXAMPLES:
    #       -> from xxx import tricky
    #       -> from astropy.io import fits
    #       -> import numpy as np
    #       -> hdu1 = fits.getdata('catalog_radec/egs_candels.fits',1)
    #       -> hdu1 = fits.getdata('catalog_radec/egs_van.fits',1)
    #       -> output=tricky.match_ctg(np.concatenate((hdu2['ra_van'].T.
    #           reshape(len(hdu2),11),hdu2['dec_van'].T.reshape(len(hdu2),1)),axis=1), 
    #           hdu1['id_candels'],np.concatenate((hdu1['ra_candels'].T.
    #           reshape(len(hdu1),11),hdu1['dec_candels'].T.reshape(len(hdu1),1)),axis=1),
    #           hdu2['id_van'])
    if target_inds is None and isinstance(target_coords,list):
        inds_target = [1]
        ra_target,dec_target = target_coords
    elif target_inds is not None:
        if isinstance(target_inds,list) or isinstance(target_inds,np.ndarray):
            inds_target = target_inds #目标源的索引
        else:
            raise TypeError('target_inds must be a list or numpy.ndarray')
        if isinstance(target_coords,np.ndarray):
            if target_coords.shape[0] != target_inds.size:
                raise ValueError('Input target_coords must have the same size as target_inds!')
            else:
                ra_target = target_coords[:,0] #目标源的坐标
                dec_target = target_coords[:,1] #目标源的坐标
    if isinstance(candidates_inds,list) or isinstance(candidates_inds,np.ndarray):
        inds_cands = candidates_inds.copy() #被匹配的源的索引
    else:
        raise TypeError('candidates_inds must be a list or numpy.ndarray')
    if candidates_coords.shape[0] != candidates_inds.size:
        print(candidates_coords.shape,candidates_inds.size)
        raise ValueError('Input candidates_coords must have the same size as candidates_inds!')
    else:
        ra_cands = candidates_coords[:,0] #被匹配的源的坐标
        dec_cands = candidates_coords[:,1] #被匹配的源的坐标

    output_inds_target_list = np.full_like
    output_inds_list = np.empty((0,2),dtype=np.int64)
    output_coords_list = np.empty((0,5),dtype=np.float64)
    
    # 建立一一对应的 空间角距离 矩阵
    coords_matrix = np.full((len(target_coords), len(candidates_inds), 4), -99.)
    coords_matrix[:,:,2], coords_matrix[:,:,0] = np.meshgrid(candidates_coords[:,0], target_coords[:,0])
    coords_matrix[:,:,3], coords_matrix[:,:,1] = np.meshgrid(candidates_coords[:,1], target_coords[:,1])
    cosarc_matrix = np.cos((coords_matrix[:,:,2] - coords_matrix[:,:,0]) * np.pi / 180.) * \
        np.cos(coords_matrix[:,:,3] * np.pi / 180.) * np.cos(coords_matrix[:,:,1] * np.pi / 180.) + \
        np.sin(coords_matrix[:,:,1] * np.pi / 180.) * np.sin(coords_matrix[:,:,3] * np.pi / 180.)
    sep_matrix = np.arccos(cosarc_matrix) * 180 / np.pi * 3600.
    
    # 按照匹配类型输出表格（矩阵算法）
    loop_flag = True
    loop_count = 0
    while loop_flag:
        target_iscloset_matrix = np.full_like(sep_matrix, fill_value=0) # 对每个 target 来说最近的 candidate 的位置
        candidate_iscloset_matrix = np.full_like(sep_matrix, fill_value=0) # 对每个 candidate 来说最近的 target 的位置
        target_iscloset_matrix[np.arange(len(target_coords)), np.argmin(sep_matrix, axis=1)] = 1
        candidate_iscloset_matrix[np.argmin(sep_matrix, axis=0), np.arange(len(candidates_coords))] = 1
        # 找到两个矩阵的交集，也就是一对一能匹配上的，先输出
        result_matrix = np.logical_and(target_iscloset_matrix, candidate_iscloset_matrix)
        indx_is_matched = np.where(result_matrix==1)
        if loop_count == 0:
            output_inds_target_list = target_inds[indx_is_matched[0]]
            output_inds_cands_list = candidates_inds[indx_is_matched[1]]
            output_ra_target_list = target_coords[:,0][indx_is_matched[0]]
            output_dec_target_list = target_coords[:,1][indx_is_matched[0]]
            output_ra_cands_list = candidates_coords[:,0][indx_is_matched[1]]
            output_dec_cands_list = candidates_coords[:,1][indx_is_matched[1]]
            output_sep_list = sep_matrix[indx_is_matched]
        else:
            output_inds_target_list = np.append(output_inds_target_list, target_inds[indx_is_matched[0]])
            output_inds_cands_list = np.append(output_inds_cands_list, candidates_inds[indx_is_matched[1]])
            output_ra_target_list = np.append(output_ra_target_list, target_coords[:,0][indx_is_matched[0]])
            output_dec_target_list = np.append(output_dec_target_list, target_coords[:,1][indx_is_matched[0]])
            output_ra_cands_list = np.append(output_ra_cands_list, candidates_coords[:,0][indx_is_matched[1]])
            output_dec_cands_list = np.append(output_dec_cands_list, candidates_coords[:,1][indx_is_matched[1]])
            output_sep_list = np.append(output_sep_list, sep_matrix[indx_is_matched])
        target_coords = np.delete(target_coords, indx_is_matched[0], axis=0)
        target_inds = np.delete(target_inds, indx_is_matched[0])
        candidates_coords = np.delete(candidates_coords, indx_is_matched[1], axis=0)
        candidates_inds = np.delete(candidates_inds, indx_is_matched[1])
        sep_matrix = np.delete(sep_matrix, indx_is_matched[0], axis=0)
        sep_matrix = np.delete(sep_matrix, indx_is_matched[1], axis=1)
        if sep_matrix.shape[0] == 0 or sep_matrix.shape[1] == 0:
            loop_flag = False
        loop_count += 1
    
    output = {'target indices':output_inds_target_list,
    'matched indices':output_inds_cands_list,
    'ra_cands':output_ra_cands_list,
    'dec_cands':output_dec_cands_list,
    'ra_targs':output_ra_target_list,
    'dec_targs':output_dec_target_list,
    'separation':output_sep_list
    }
    
    return output

# 把id,ra,dec转换成ds9的region文件
def coords2reg(id, ra, dec, size=None, color=None, text=True):
    # size unit: "
    template = '''# Region file format: DS9 version 4.1
    global color=green dashlist=8 3 width=1 font="helvetica 10 normal roman" select=1 highlite=1 dash=0 fixed=0 edit=1 move=1 delete=1 include=1 source=1
    fk5
    '''
                    # circle(215.0566198,52.9277772,25.558")
                    # circle(214.9304338,52.85543778,25.558") # color=red'''
                    # # text(215.14186,52.967176) text={27}
    for num in range(len(id)):
        if size == None:
            string_size = '1"'
        elif len(size) == 1:
            string_size = str(size[0])+'"'
        else:
            string_size = str(size[num])+'"'
        if color == None:
            string_color = ''
        elif len(color)==1:
            string_color = color[0]
        else:
            string_color = '# color='+color[num]
        if text:
            string = 'circle('+str(ra[num])+','+str(dec[num])+','+string_size+') '+string_color+' \n# text('+str(ra[num])+','+str(dec[num])+') text={'+str(id[num])+'} \n'
        else:
            string = 'circle('+str(ra[num])+','+str(dec[num])+','+string_size+') '+string_color+' \n'
        if num == 0:
            region_text = template + string
        else:
            region_text += string
    return region_text

def filtermag(wavemodel,fluxmodel,filterwave,filtertrans):
    from scipy.interpolate import interp1d
    from scipy.integrate import simpson
    import scipy.constants as cst
    # wavemodel unit: um
    # fluxmodel unit: Jy
    nu_model = cst.c*1.e2 / (wavemodel*1.e-4) # unit:Hz
    fnu_model = fluxmodel.copy()
    nu_filter = cst.c*1.e2 / (filterwave*1.e-4)
    numerator = (simpson(interp1d(nu_model,fnu_model)(nu_filter)/nu_filter*filtertrans, nu_filter))
    denominator = (integrate.simpson(3631/nu_filter*filtertrans, nu_filter))
    magnitude = -2.5*np.log10(numerator/denominator)
    return magnitude

# Calulate number counts
def numbercounts(mag, magbinstep=None, logbin=False, logbincounts=None):
    maglist = mag.copy()
    ranlow, ranhigh = [np.min(mag), np.max(mag)]
    if logbin==False and magbinstep==None:
        magbinstep = 0.05
    if logbincounts==None:
        logbinscounts=11
    if logbin:
        magbins = np.logspace(np.log10(ranlow), np.log10(ranhigh), logbincounts, base=10.)
    else:
        magbins = np.arange(ranlow, ranhigh, magbinstep)
    counts = np.zeros(len(magbins)-1)
    for num in range(len(magbins)-1):
        if num==len(magbins)-1:
            maglistinrange = maglist[(maglist<=magbins[num+1])&(maglist>=magbins[num])]
        else:
            maglistinrange = maglist[(maglist<magbins[num+1])&(maglist>=magbins[num])]
        counts[num] = len(maglistinrange)
    return magbins, counts

# define a class, Point, with attributions x, y (coordinates)
class Point:
    def __init__(self, x, y):
        self.x = x
        self.y = y

# Checking if a point is inside a polygon
# This function can only apply to the polygon without intersected empty region
def point_in_polygon(pointcoords, polygonpoints):
    polygon = []
    for borderpoint in polygonpoints:
        polygon.append(Point(borderpoint[0],borderpoint[1]))

    num_border = len(polygon)
    point = Point(pointcoords[0],pointcoords[1])
    x, y = point.x, point.y
    inside = False

    # Store the first point in the polygon and initialize the second point
    p1 = polygon[0]

    # Loop through each edge in the polygon
    for i in range(1, num_border):
        # Get the next point in the poly gon
        p2 = polygon[i]

        # Check if the point is above the minimum y coordinate of the edge
        if y > min(p1.y, p2.y):
            # Check if the point is below the maximum y coordinate of the edge
            if y <= max(p1.y, p2.y):
                # Check if the point is to the left of the maximum x coordinate of the edge
                if x <= max(p1.x, p2.x):
                    # Calculate the x-intersection of the line connecting the point to the edge
                    x_intersection = (y - p1.y) * (p2.x - p1.x) / (p2.y - p1.y) + p1.x

                    # Check if the point is on the same line as the edge or to the left of the x-intersection
                    if p1.x == p2.x or x <= x_intersection:
                        # Flip the inside flag
                        inside = not inside

        # Store the current point as the first point for the next iteration
        p1 = p2

    # Return the value of the inside flag
    return inside

# determine the border according to a catalog with ra,dec
def border(ra,dec, rabins=None):
    ramin,ramax = [np.min(ra),np.max(ra)]
    if rabins==None:
        rabins = 31
    decedges = np.zeros((rabins-1,2))
    raedges = np.zeros((rabins-1,2))
    rabinedges = np.linspace(ramin,ramax, rabins)
    for num in range(rabins-1):
        ralow, rahigh = [rabinedges[num],rabinedges[num+1]]
        if num==rabins-2:
            decbin = dec[(ra>=ralow)&(ra<=rahigh)]
            rabin = ra[(ra>=ralow)&(ra<=rahigh)]
        else:
            decbin = dec[(ra>=ralow)&(ra<rahigh)]
            rabin = ra[(ra>=ralow)&(ra<rahigh)]
        if decbin.size==0:
            declow,dechigh = [-99.,-99.]
            radeclow,radechigh = [-99.,-99.]
        else:
            declow, dechigh = [np.min(decbin), np.max(decbin)]
            radeclow, radechigh = [rabin[np.argmin(decbin)], rabin[np.argmax(decbin)]]
        decedges[num] = [declow, dechigh]
        raedges[num] = [radeclow, radechigh]
        
    borderarray = np.zeros((rabins+1+rabins,2))

    borderarray[0] = [ra[np.argmin(ra)],dec[np.argmin(ra)]]
    borderarray[rabins] = [ra[np.argmax(ra)],dec[np.argmax(ra)]]

    borderarray[1:rabins,:] = np.concatenate((raedges[:,0].reshape(rabins-1,1), (decedges[:,0]).reshape(rabins-1,1)),axis=1)
    borderarray[rabins+1:-1,:] = np.concatenate((raedges[:,1].reshape(rabins-1,1)[::-1], (decedges[:,1]).reshape(rabins-1,1)[::-1]),axis=1)
    # The first and the last one are the same point, which means this border is a closed polygon
    borderarray[-1] = [ra[np.argmin(ra)],dec[np.argmin(ra)]]

    return borderarray